import base64
import io
import re
from typing import Any, Dict, Optional

import numpy as np
import paddleocr
from flask import Flask, jsonify, request
from PIL import Image

app = Flask(__name__)

# 身份证正则表达式
ID_CARD_REGEX = (
    "("
    r"([1-9]\d{5}\d{2}((0[1-9])|(10|11|12))(([0-2][1-9])|10|20|30|31)\d{3})"
    "|"
    r"([1-9]\d{5}(18|19|([23]\d))\d{2}((0[1-9])|(10|11|12))(([0-2][1-9])|10|20|30|31)\d{3}[0-9Xx])"
    ")"
)

# 全局 ppocr 对象
ocr_engine = paddleocr.PaddleOCR(
    lang="ch",
    use_angle_cls=True,
    use_gpu=False,
    show_log=False,
)


def match_idcard(text: str) -> Optional[str]:
    """匹配身份证

    Args:
        text (str): 输入文本

    Returns:
        str: 身份证, 如果匹配不到返回 None
    """
    m = re.search(ID_CARD_REGEX, text)
    return m.group(1) if m is not None else None


def base64_to_ndarray(data: Dict[Any, Any]) -> np.ndarray:
    """解析json输入, 并将图片从原来 base64格式 转为 ndarray

    Args:
        data (Dict[Any, Any]): json 输入

    Returns:
        np.NDArray[Any]: ndarray 格式的图片
    """
    imgbyte = base64.b64decode(data["image"])
    image = io.BytesIO(imgbyte)
    temp = Image.open(image)
    img_ndarray = np.array(temp)[:, :, :3]
    return img_ndarray


@app.errorhandler(Exception)
def handle_exception(e):
    return jsonify({"code": 500, "result": "", "message": repr(e)}), 500


@app.route("/idcard-ocr", methods=["POST"])
def idcard_ocr():
    img = base64_to_ndarray(request.json)
    result = ocr_engine.ocr(img)
    for res in result:
        for line in res:
            _, (text, _) = line
            result = match_idcard(text.strip())
            if result:
                return jsonify({"code": 200, "result": result})

    return jsonify({"code": 400, "result": "", "message": "OCR failed"})


def main():
    app.run(debug=True, host="127.0.0.1", port=5000)


if __name__ == "__main__":
    main()
